#lang racket

(require "sets.rkt")
(require "syntax.rkt")

(provide scope binding-table environment
         top-scope
         annotate/expand serialize-syntax
         eval)

;;;;;;;;;;;;

(define (assoc-cdr key table)
  (cond ((assoc key table) => (lambda (result) (cdr result)))
        (else #f)))

(define (update alist k v)
  (let aux ((kvs alist) (out (lambda (x) x)))
    (cond ((null? kvs) (out (list (cons k v))))
          ((eq? k (caar kvs)) (out (cons (cons k v) (cdr kvs))))
          (else (aux (cdr kvs) (lambda (rest) (out (cons (car kvs) rest))))))))

;;;;;;;;;;;;

(define (lookup-binding-table env name scope)
  (let ((scope-table (hash-ref env name #f)))
    (if scope-table
        (let loop ((mappings (hash->list scope-table)) (top (cons -1 #f)))
          (if (null? mappings)
              (if (= 0 (car top))
                  (error "unbound variable" name)
                  (cdr top))
              (let* ((m (car mappings))
                     (m-scope (car m))
                     (m-val (cdr m))
                     (score (subset-score scope m-scope)))
                ;; (displayln `(subset-score ,(set->list scope) ,(set->list m-scope)
                ;;                           = ,score = ,name => ,m-val))
                (if (> score (car top))
                    (loop (cdr mappings) (cons score m-val))
                    (loop (cdr mappings) top)))))
        #f)))

(define (extend-binding-table env name scope value)
  (let* ((inner-env (hash-ref env name (hash)))
         (inner-env^ (hash-set inner-env scope value)))
    (hash-set env name inner-env^)))

;;;;;;;;;;;

(define empty-scope 0)
(define empty-binding-table (hash))
(define empty-environment '())
(define empty-definition-context (list (set) '() '()))

(define scope (make-parameter empty-scope))
(define binding-table (make-parameter empty-binding-table))
(define environment (make-parameter empty-environment))
(define definition-context (make-parameter empty-definition-context))

(define (next-scope!)
  (let ((s (+ 1 (scope))))
    (scope s)
    s))
  
(define (add-binding! name scope expander)
  (let ((name* (gensym name)))
    (binding-table (extend-binding-table (binding-table) name scope name*))
    (environment `((,name* . ,expander) . ,(environment)))
    name*))

(define (add-use-site-scope! scope)
  (let ((def-ctx (definition-context)))
    (definition-context (list (set-union (car def-ctx) scope)
                              (cadr def-ctx)
                              (caddr def-ctx)))))

(define (add-syntax-definition! name value)
  (let ((def-ctx (definition-context)))
    (definition-context (list (car def-ctx)
                              (cons (cons name value) (cadr def-ctx))
                              (caddr def-ctx)))))

(define (add-value-definition! name value)
  (let ((def-ctx (definition-context)))
    (definition-context (list (car def-ctx)
                              (cadr def-ctx)
                              (cons (cons name value) (caddr def-ctx))))))

(define (stop! sy)
  (let ((binding (resolve-binding sy))
        (expander (expander-for sy)))
    (environment (update (environment) binding (lambda (s) s)))
    (cons sy expander)))

(define (unstop! sy expander)
  (let ((binding (resolve-binding sy)))
    (environment (update (environment) binding expander))))

;;;;;;;;;;;;;

(define (resolve-binding sy)
  (cond ((syntax-identifier? sy)
         (lookup-binding-table (binding-table) (syntax-expression sy) (syntax-scope sy)))
        (else #f)))

(define (free-identifier=? a b)
  (let ((a* (resolve-binding a))
        (b* (resolve-binding b)))
    (and a* b* (eq? a* b*))))

(define (bound-identifier=? a b)
  (equal? (set->list (syntax-scope a)) (set->list (syntax-scope b))))

(define (expander-for sy)
  (cond ((resolve-binding sy) => (lambda (b) (assoc-cdr b (environment))))
        (else #f)))

;;;;;;;;;;;;;;;;;

(define (annotate/expand sy)
  (cond ((syntax-pair? sy)
         (let* ((head (syntax-car sy)))
           (cond ((expander-for head) => (lambda (expand) (expand sy)))
                 (else (syntax-map annotate/expand sy)))))
        (else sy)))

(define (annotate/expand* sy stop-list)
  (let ((stopped (map stop! stop-list)))
    (let ((v (annotate/expand sy)))
      (for-each (lambda (sy-expander) (unstop! (car sy-expander) (cdr sy-expander))) stopped)
      v)))

(define (serialize-syntax sy)
  (cond ((syntax-pair? sy)
         (let ((e (syntax-expression sy)))
           (if (eq? (resolve-binding (car e)) syntax-gensym)
               (list (serialize-syntax (car e)) (syntax-car (cdr e)))
               (cons (serialize-syntax (car e)) (serialize-syntax (cdr e))))))
        ((syntax-atom? sy) (syntax-expression sy))
        ((syntax-identifier? sy) (or (resolve-binding sy) (syntax-expression sy)))
        (else (error "not a syntax object" sy))))

(define (make-macro-expander sy)
  (let ((expander-fn (eval (serialize-syntax sy))))
    (lambda (sy)
      (let ((use-site-scope (set (next-scope!))))
        (add-use-site-scope! use-site-scope)
        (annotate/expand (syntax-add-scope (expander-fn (eval sy)) use-site-scope))))))

;;;;;;;;

(define (expand-syntax sy) sy)

(define (expand-quasisyntax sy)
  (define (quasisyntax sy n)
    (if (syntax-pair? sy)
        (cond ((eq? 'unsyntax (syntax-expression (syntax-car sy)))
               (if (= n 1)
                   (syntax-car (syntax-cdr sy))
                   (datum->syntax
                    `(syntax-cons
                      (syntax ,(syntax-car sy))
                      (syntax-cons ,(quasisyntax (syntax-car (syntax-cdr sy)) (- n 1))
                                   ())))))
              ((eq? 'quasisyntax (syntax-expression (syntax-car sy)))
               (if (= n 0)
                   (quasisyntax (syntax-car (syntax-cdr sy)) (+ n 1))
                   (datum->syntax
                    `(syntax-cons
                      (syntax ,(syntax-car sy))
                      (syntax-cons ,(quasisyntax (syntax-car (syntax-cdr sy)) (+ n 1))
                                   ())))))
              (else
               (datum->syntax
                `(syntax-cons ,(quasisyntax (syntax-car sy) n)
                              ,(quasisyntax (syntax-cdr sy) n)))))
        (datum->syntax (list 'syntax sy))))
  (annotate/expand (syntax-add-scope (quasisyntax sy 0) top-scope)))

(define (expand-begin sy)
  (syntax-cons (syntax-car sy) (syntax-map annotate/expand (syntax-cdr sy))))

(define (expand-if sy)
  (syntax-cons (syntax-car sy) (syntax-map annotate/expand (syntax-cdr sy))))

(define (expand-body sy)
  ;; fresh definition-context
  (parameterize ((definition-context empty-definition-context))
    ;; partial expand body, collecting definitions
    (let* ((body-scope (set (next-scope!)))
           (body (syntax-add-scope sy body-scope))
           (body^ (annotate/expand* body core-stop-list))
           (def-ctx (definition-context))
           (use-site-scopes (car def-ctx))
           (syntax-definitions (cadr def-ctx))
           (definitions (caddr def-ctx))
           (syntax-definitions^
            (map (lambda (d)
                   `(,(syntax-remove-scopes (car d) use-site-scopes) ,(cdr d)))
                 syntax-definitions))
           (definitions^
             (map (lambda (d)
                    `(,(syntax-remove-scopes (car d) use-site-scopes) ,(cdr d)))
                  definitions))
           (template body^)
           (template^ (if (null? definitions^) template `(letrec ,definitions^ ,template)))
           (template^^ (if (null? syntax-definitions^)
                           template^
                           `(letrec-syntax ,syntax-definitions^ ,template^))))
      (annotate/expand (syntax-add-scope (datum->syntax template^^) top-scope)))))

(define (expand-define sy)
  (let ((id (syntax-car (syntax-cdr sy)))
        (value (syntax-car (syntax-cdr (syntax-cdr sy)))))
    (add-value-definition! id value)
    (datum->syntax `(begin))))

(define (expand-define-syntax sy)
  (let ((id (syntax-car (syntax-cdr sy)))
        (value (syntax-car (syntax-cdr (syntax-cdr sy)))))
    (add-syntax-definition! id value)
    (datum->syntax `(begin))))

(define (expand-lambda sy)
  (let ((params (syntax-car (syntax-cdr sy)))
        (body (syntax-car (syntax-cdr (syntax-cdr sy))))
        (scope (set (next-scope!))))
    (let ((params* (syntax-add-scope params scope)))
      (syntax-map (lambda (p) (add-binding! (syntax-expression p) (syntax-scope p) #f))
                  params*)
      (datum->syntax
       `(,(syntax-car sy) ,params*
         ,(expand-body (syntax-add-scope body scope)))))))

(define (expand-let-syntax sy)
  (let ((bindings (syntax-car (syntax-cdr sy)))
        (body (syntax-car (syntax-cdr (syntax-cdr sy))))
        (scope (set (next-scope!))))
    (syntax-map
     (lambda (binding)
       (let ((binder (syntax-add-scope (syntax-car binding) scope))
             (rhs (annotate/expand (syntax-car (syntax-cdr binding)))))
         (add-binding! (syntax-expression binder) (syntax-scope binder)
                       (make-macro-expander rhs))))
     bindings)
    (expand-body (syntax-add-scope body scope))))

(define (expand-letrec-syntax sy)
  (let ((bindings (syntax-car (syntax-cdr sy)))
        (body (syntax-car (syntax-cdr (syntax-cdr sy))))
        (scope (set (next-scope!))))
    (syntax-map
     (lambda (binding)
       (let ((binder (syntax-add-scope (syntax-car binding) scope))
             (rhs (annotate/expand (syntax-add-scope (syntax-car (syntax-cdr binding)) scope))))
         (add-binding! (syntax-expression binder) (syntax-scope binder)
                       (make-macro-expander rhs))))
     bindings)
    (expand-body (syntax-add-scope body scope))))

(define (expand-let sy)
  (let ((bindings (syntax-car (syntax-cdr sy)))
        (body (syntax-car (syntax-cdr (syntax-cdr sy))))
        (scope (set (next-scope!))))
    (let ((bindings* (syntax-map
                      (lambda (binding)
                        (let ((binder (syntax-add-scope (syntax-car binding) scope))
                              (rhs (annotate/expand (syntax-car (syntax-cdr binding)))))
                          (add-binding! (syntax-expression binder) (syntax-scope binder)
                                        #f)
                          (datum->syntax (list binder rhs))))
                      bindings)))
      (datum->syntax
       `(,(syntax-car sy) ,bindings*
         ,(expand-body (syntax-add-scope body scope)))))))

(define (expand-letrec sy)
  (let ((bindings (syntax-car (syntax-cdr sy)))
        (body (syntax-car (syntax-cdr (syntax-cdr sy))))
        (scope (set (next-scope!))))
    (let ((bindings* (syntax-map
                      (lambda (binding)
                        (let ((binder (syntax-add-scope (syntax-car binding) scope))
                              (rhs (annotate/expand (syntax-add-scope (syntax-car (syntax-cdr binding)) scope))))
                          (add-binding! (syntax-expression binder) (syntax-scope binder)
                                        #f)
                          (datum->syntax (list binder rhs))))
                      bindings)))
      (datum->syntax
       `(,(syntax-car sy) ,bindings*
         ,(expand-body (syntax-add-scope body scope)))))))

(define top-scope (set (next-scope!)))
(define begin-gensym (add-binding! 'begin top-scope expand-begin))
(define syntax-gensym (add-binding! 'syntax top-scope expand-syntax))
(define quasisyntax-gensym (add-binding! 'quasisyntax top-scope expand-quasisyntax))
(define if-gensym (add-binding! 'if top-scope expand-if))
(define lambda-gensym (add-binding! 'lambda top-scope expand-lambda))
(define let-syntax-gensym (add-binding! 'let-syntax top-scope expand-let-syntax))
(define letrec-syntax-gensym (add-binding! 'letrec-syntax top-scope expand-letrec-syntax))
(define let-gensym (add-binding! 'let top-scope expand-let))
(define letrec-gensym (add-binding! 'letrec top-scope expand-letrec))
(define define-gensym (add-binding! 'define top-scope expand-define))
(define define-syntax-gensym (add-binding! 'define-syntax top-scope expand-define-syntax))

(define stop-gensym (add-binding! 'stop-gensym top-scope (lambda (sy) sy)))

(define core-stop-list
  (map (lambda (s) (syntax-add-scope (datum->syntax s) top-scope))
       '(lambda let letrec let-syntax letrec-syntax)))

;;;;;;;;;;;;;;

(define (eval0 exp env)
  ;;(displayln exp)
  (match exp
    [(? null?) exp]
    [(? syntax?) exp]
    [(? number?) exp]
    [(? boolean?) exp]
    [(? string?) exp]
    [(? symbol?) (cond ((assoc exp (unbox env)) => cdr)
                       (else (error "unbound" exp)))]
    [`(,(== begin-gensym) . ,body)
     (let loop ((xs body) (last '()))
       (if (null? xs)
           last
           (loop (cdr xs) (eval0 (car xs) env))))]
    [`(,(== 'set!) ,a ,v) (set-box! env (update (unbox env) (eval0 a env) (eval0 v env)))]
    [`(,(== if-gensym) ,t ,c ,a) (if (eval0 t env) (eval0 c env) (eval0 a env))]
    [`(,(== syntax-gensym) ,stx) stx]
    [`(,(== lambda-gensym) ,params ,body)
     (lambda args
       (eval0 body (box (append (map cons params args) (unbox env)))))]
    [`(,(== let-gensym) ,bindings ,body)
     (eval0 body (box (append (map (lambda (b) (cons (car b) (eval0 (car (cdr b)) env)))
                                   bindings)
                              (unbox env))))]
    [`(,(== letrec-gensym) ,bindings ,body)
     (let ((env^ (append (map (lambda (b) (cons (car b) #f)) bindings) (unbox env))))
       (let ((env^^ (foldl (lambda (b env)
                             (update env (car b) (eval0 (car (cdr b)) (box env^))))
                           env^ bindings)))
         (eval0 body (box env^^))))]
    [`(,fn . ,args)
     (let ((fn^ (eval0 fn env))
           (args^ (map (lambda (e) (eval0 e env)) args)))
       (apply fn^ args^))]))

(define base-env
  `((syntax-car . ,syntax-car)
    (syntax-cdr . ,syntax-cdr)
    (syntax-null? . ,syntax-null?)
    (syntax-cons . ,syntax-cons)
    (eq? . ,eq?)
    (cons . ,cons)
    (list . ,list)
    (car . ,car)
    (cdr . ,cdr)
    (+ . ,+)
    (* . ,*)
    (display . ,display)
    (for-each . ,for-each)))

(define (eval exp) (eval0 exp (box base-env)))
